import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.List;

/**
 * <p>
 *
 * </p>
 * @author wind
 * @date    2024/12/12 14:13
 * @version v1.0
 */
public class YarnHelper {

    private static final Logger logger = LoggerFactory.getLogger(YarnHelper.class);

    private final String hadoopConfDir;

    private final String krb5User;

    private final String krb5Tab;

    private final String krb5Conf;

    private YarnHelper(String hadoopConfDir, String krb5User, String krb5Tab, String krb5Conf){
        this.hadoopConfDir = hadoopConfDir;
        this.krb5User = krb5User;
        this.krb5Tab = krb5Tab;
        this.krb5Conf = krb5Conf;
    }

    public static YarnHelper fromConfig(String hadoopConfDir){
        return new YarnHelper(hadoopConfDir,null, null, null);
    }

    public static YarnHelper fromKrb5(String hadoopConfDir, String krb5User, String krb5Tab, String krb5Conf){
        return new YarnHelper(hadoopConfDir, krb5User, krb5Tab, krb5Conf);
    }


    public Configuration initConf(){
        YarnConfiguration conf = new YarnConfiguration();
        if(StringUtils.isNotEmpty(hadoopConfDir)){
            conf.addResource(new Path(hadoopConfDir + File.separator + "core-site.xml"));
            conf.addResource(new Path(hadoopConfDir + File.separator + "hdfs-site.xml"));
            conf.addResource(new Path(hadoopConfDir + File.separator + "yarn-site.xml"));
        }
        if(isKerberos()){
            System.setProperty("java.security.krb5.conf", krb5Conf);
            try {
                UserGroupInformation.setConfiguration(conf);
                UserGroupInformation.loginUserFromKeytab(krb5User, krb5Tab);
            } catch (IOException e) {
                logger.error(e.getMessage());
            }
        }
        return conf;
    }

    /**
     * 获取客户端
     * @return
     */
    public YarnClient getClient(){
        YarnClient yarnClient = YarnClient.createYarnClient();
        yarnClient.init(initConf());
        yarnClient.start();
        return yarnClient;
    }

    /**
     * 是否配置Kerberos认证
     * @return
     */
    public boolean isKerberos(){
        return StringUtils.isNotEmpty(krb5Conf) && StringUtils.isNotEmpty(krb5User) && StringUtils.isNotEmpty(krb5Tab);
    }


    /**
     * 停止任务
     * @param appId
     * @throws Exception
     */
    public void stopJob(String appId) throws Exception {
        ApplicationId applicationId = ApplicationId.fromString(appId);
        action(client -> {
            ApplicationReport applicationReport = client.getApplicationReport(applicationId);
            if (applicationReport != null){
                switch (applicationReport.getYarnApplicationState()){
                    case NEW:
                    case NEW_SAVING:
                    case ACCEPTED:
                    case RUNNING:
                    case SUBMITTED:
                        client.killApplication(applicationId);
                        break;
                    default:
                        break;
                }
            }
        });
    }

    /**
     * 查询任务状态
     * @param appId
     * @throws Exception
     */
    public void getJobStatus(String appId) throws Exception {
        ApplicationId applicationId = ApplicationId.fromString(appId);
        action(client -> {
            ApplicationReport applicationReport = client.getApplicationReport(applicationId);
            System.out.println(applicationReport.getYarnApplicationState().toString());
        });
    }

    /**
     * 查询任务状态
     * @throws Exception
     */
    public void getJobList() throws Exception {
        action(client -> {
            List<ApplicationReport> jobs = client.getApplications();
            for(ApplicationReport job : jobs){
                System.out.println(job);
            }
        });
    }


    /**
     * hdfs操作
     * @param callback
     * @throws Exception
     */
    public void action(Callback callback) throws Exception {
        try(YarnClient client = getClient()){
            if(callback != null){
                callback.call(client);
            }
        }
    }

    @FunctionalInterface
    interface Callback{
        void call(YarnClient client) throws Exception;
    }
}
